#!/usr/bin/env python3
# run_scan.py
import math
import os
import pathlib as pl
import subprocess
import sys
import threading
from string import Template
from genfigerprint import genfingerprint
# from icecream import ic

##0000000000000000000000000000000000000000 用户可调参数
##================================ Geant4 exe 文件路径
_CWD = pl.Path(__file__).parent
_G4_PRJ_DIR = _CWD.parent / "G4"
_G4_BUILD_TOP = _G4_PRJ_DIR / "build"
_G4DIR = _G4_BUILD_TOP / "Release"
_G4EXE = _G4DIR / "g4test"
_G4THREADS = 20
## em, em_op3, em_op4, QBBC, QBBC_IBC
_G4PHYLIST = "em_op4"
_G4BEAMON = 10000

# G4_EXE = './g4test'         # Linux/macOS 例
##=====================================
_NPOINTS = 20  # 采样点数
_EMIN, _EMAX = 0.1, 500.0  # 能量范围 / MeV
_LOG_SCALE = True  # True=对数采样, False=线性采样
_MAC_DIR = _CWD / "macros"  # Geant4 input mac files
_OUT_DIR = _CWD / "out"  # 所有输出文件放在这个文件夹
_LOG_DIR = _MAC_DIR  # log files

##0000000000000000000000000000000000000000
# 仅在当前进程(及它的所有子进程)里临时追加 PATH
_fingerprint = genfingerprint()
## 准备要追加到 PATH 的目录
_HOMEDIR = pl.Path.home()
_G4BIN_DIR = ""
_G4DATA_DIR = ""
## ymy linux
if _fingerprint == "4eac38846a626f7a001d86647223149d":
    _G4BIN_DIR = pl.Path(_HOMEDIR / "cpplibs/geant4v11.3.2/bin").resolve()
    _G4DATA_DIR = pl.Path(_HOMEDIR / "cpplibs/g4dataset11.3.2").resolve()
## ymy des fedora gnome
elif _fingerprint == "21c130c2e0a73c2b05e5af6bbe8aa042":
    _G4BIN_DIR = pl.Path(r"C:/cpplibs/geant4-11.3.2/bin").resolve()
    _G4DATA_DIR = pl.Path(_HOMEDIR / "cpplibs/G4Data11.3.2").resolve()

os.environ["PATH"] = f"{_G4BIN_DIR}{os.pathsep}{os.environ.get('PATH', '')}"
# ic(os.environ["PATH"])
os.environ["GEANT4_DATA_DIR"] = _G4DATA_DIR.as_posix()


##0000000000000000000000000000000000000000  确保输出目录存在
os.makedirs(_OUT_DIR, exist_ok=True)
os.makedirs(_MAC_DIR, exist_ok=True)
os.makedirs(_LOG_DIR, exist_ok=True)

##000000000000000000000000000000000000000 字符串模板(与题目给出的 letd.mac 完全一致, 仅把能量处留 $energy)
mac_template = Template(r"""##======================== 全局开关与线程数
/control/verbose 2
/tracking/verbose 0
/run/verbose 2
/event/verbose 0
## 修改默认线程数, 必须在 /run/initialize前
/run/numberOfThreads ${g4threads}

##======================== physics list
/testem/phys/addPhysics ${g4phylist}

##======================== 探测器几何, water phantom
/testem/det/verbose 0
/testem/det/PhantomXY 300 mm
/testem/det/PhantomZ 300 mm
## world= Factor*Phantom
/testem/det/Factor 45

##======================== 初始化, 之后不能再改几何/物理/线程
/run/initialize

##========================
/run/setCut 0.05 mm

##========================  定义 scorer/mesh, 必须在 beamOn 之前
/score/create/boxMesh boxMesh1
/score/mesh/boxSize 150.0 150.0 150.0 mm
/score/mesh/nBin 1 1 150
/score/quantity/energyDeposit eDep
/letd/addToMesh boxMesh1 LETd
/score/close

/score/create/boxMesh boxMesh2
/score/mesh/boxSize 150.0 150.0 150.0 mm
/score/mesh/nBin 1 1 1
/rgProj/addToMesh boxMesh2 rgProj1
/score/close

/score/list

##======================== 粒子源, 必须在 beamOn 之前
/gun/particle ion
/gun/ion 2 4
/testem/gun/beamE ${g4energy} MeV
##---------------- !!!careful of the location of gun setting
/testem/gun/beamZ -6632 mm
#/testem/gun/spotDeltaY 10 mm
#/testem/gun/sigmaX 4 mm
#/testem/gun/sigmaY 4 mm
#/testem/gun/distypeXY gauss

##======================== run event
/run/beamOn ${g4beamon}

##======================== Dump scorers to files; beamOn 之后
/score/dumpQuantityToFile boxMesh1  eDep eDep_${g4energy}MeV.csv
/score/dumpQuantityToFile boxMesh1  LETd LETd_${g4energy}MeV.csv
/score/dumpQuantityToFile boxMesh2  rgProj1 rgProj_${g4energy}MeV.csv
""")


def generate_energy_list(n, emin, emax, log_scale=True):
    if False:
        """生成能量采样列表(单位 MeV)"""
        if log_scale:
            logmin, logmax = math.log10(emin), math.log10(emax)
            delta = (logmax - logmin) / (n - 1)
            return [10 ** (logmin + i * delta) for i in range(n)]
        else:
            delta = (emax - emin) / (n - 1)
            return [emin + i * delta for i in range(n)]
    return [
        7.000e02,
        1.000e03,
    ]
    return [
        1.000e-03,
        5.000e-03,
        1.000e-02,
        5.000e-02,
        1.000e-01,
        5.000e-01,
        1.000e00,
        5.000e00,
        1.000e01,
        5.000e01,
        1.000e02,
        5.000e02,
        7.000e02,
        1.000e03,
    ]


##===============================
def format_energy(energy):
    """把 0.1/12.3/500.0 之类格式化成小数点后 1 位, 方便文件名统一"""
    return f"{energy:.1f}"


##===============================
def tee_thread(pipe, file_obj):
    """把 pipe 中的内容同时写到文件与终端"""
    for line in iter(pipe.readline, b""):
        line_dec = line.decode(errors="ignore")
        file_obj.write(line_dec)
        file_obj.flush()
        sys.stdout.write(line_dec)
        sys.stdout.flush()
    pipe.close()


##===============================
def run_with_tee(cmd_list, log_path):
    """
    纯 Python 实现: 运行 cmd_list, 把 stdout+stderr 实时打印并写入 log_path
    cmd_list: 可执行文件路径及参数, 如 ['.\\g4test', 'letd_10.0MeV.mac']
    """
    with open(log_path, "w", encoding="utf-8") as log_file:
        ## 启动子进程
        proc = subprocess.Popen(
            cmd_list, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=1
        )
        ## 启动后台线程做 tee
        thread = threading.Thread(
            target=tee_thread, args=(proc.stdout, log_file), daemon=True
        )
        thread.start()
        ## 等待进程结束
        ret = proc.wait()
        thread.join()
        return ret


##===============================
def main():
    os.chdir(_OUT_DIR)
    energies = generate_energy_list(_NPOINTS, _EMIN, _EMAX, _LOG_SCALE)
    print(f"即将扫描 {len(energies)} 个能量点: ")
    print([format_energy(e) for e in energies])

    for engy in energies:
        estr = format_energy(engy)
        mac_file = _MAC_DIR / f"letd_{estr}MeV.mac"
        log_file = _LOG_DIR / f"log_{estr}MeV.log"
        ## 写宏文件
        with open(mac_file, "w", encoding="utf-8") as f:
            f.write(
                mac_template.substitute(
                    g4threads=_G4THREADS,
                    g4phylist=_G4PHYLIST,
                    g4energy=estr,
                    g4beamon=_G4BEAMON,
                )
            )
        ## cmd_list 直接就是 [G4_EXE, mac_file]
        cmd_list = [_G4EXE, mac_file]
        # ic(cmd_list)
        # continue
        ret = run_with_tee(cmd_list, log_file)
        if ret != 0:
            print(f"!!! 能量 {estr} MeV 运行失败, returncode = {ret}")
    os.chdir(_CWD)
    print("\n全部扫描完成! 结果保存在目录: ", _OUT_DIR)


if __name__ == "__main__":
    main()
